#!/bin/bash
# 模型导出脚本 - 将YOLO PT模型导出为TensorRT引擎格式
# 使用conda环境: dl_env

# 激活conda环境
source /home/lixinlong/SoftWare/miniconda3/etc/profile.d/conda.sh
conda activate dl_env

# 设置工作目录
# 兼容bash和zsh
if [ -n "${BASH_SOURCE[0]}" ]; then
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
else
    SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
fi
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"

# 默认参数
PT_PATH="/home/lixinlong/Project/linden_belt_detection/src/linden_belt_detection/model/pt/best.pt"
IMGSZ=640
DEVICE=0
HALF=true
BATCH=1
NMS=true
DATA="/home/lixinlong/Project/yolo-train/dataset.yaml"
DYNAMIC=false
INT8=false
FRACTION=1.0
WORKSPACE=8
CONF=0.5
IOU=0.5
MAX_DET=50
QUIET=false

# 打印使用说明
usage() {
    cat << EOF
使用方法: $0 --pt_path <模型路径> [选项]

必需参数:
  --pt_path PATH          PT模型文件路径

可选参数:
  --imgsz SIZE            输入图像尺寸 (默认: 640)
  --device ID             GPU设备ID (默认: 0)
  --no-half               禁用FP16半精度模式
  --batch SIZE            批处理大小 (默认: 1)
  --nms                   将NMS融合到导出的模型中
  --data PATH             数据集配置文件路径，用于INT8校准
  --dynamic               启用动态形状支持
  --int8                  使用INT8量化
  --fraction VALUE        用于INT8校准的数据集比例 (默认: 1.0)
  --workspace SIZE        TensorRT工作空间大小(GiB) (默认: 8)
  --conf VALUE            NMS置信度阈值 (默认: 0.25)
  --iou VALUE             NMS IOU阈值 (默认: 0.45)
  --max-det NUM           最大检测数量 (默认: 100)
  --quiet                 静默模式
  -h, --help              显示此帮助信息

示例:
  # 基本导出
  $0 --pt_path output/minion_detect/weights/best.pt

  # 使用FP16和动态形状
  $0 --pt_path model.pt --dynamic --workspace 8

  # 使用INT8量化
  $0 --pt_path model.pt --int8 --data dataset.yaml --fraction 0.5

  # 融合NMS并使用自定义阈值
  $0 --pt_path model.pt --nms --conf 0.3 --iou 0.5 --max-det 200

  # 完整配置
  $0 --pt_path model.pt --imgsz 640 --batch 4 --nms --dynamic --workspace 8 --conf 0.3

EOF
    exit 1
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --pt_path)
            PT_PATH="$2"
            shift 2
            ;;
        --imgsz)
            IMGSZ="$2"
            shift 2
            ;;
        --device)
            DEVICE="$2"
            shift 2
            ;;
        --no-half)
            HALF=false
            shift
            ;;
        --batch)
            BATCH="$2"
            shift 2
            ;;
        --nms)
            NMS=true
            shift
            ;;
        --data)
            DATA="$2"
            shift 2
            ;;
        --dynamic)
            DYNAMIC=true
            shift
            ;;
        --int8)
            INT8=true
            shift
            ;;
        --fraction)
            FRACTION="$2"
            shift 2
            ;;
        --workspace)
            WORKSPACE="$2"
            shift 2
            ;;
        --conf)
            CONF="$2"
            shift 2
            ;;
        --iou)
            IOU="$2"
            shift 2
            ;;
        --max-det)
            MAX_DET="$2"
            shift 2
            ;;
        --quiet)
            QUIET=true
            shift
            ;;
        -h|--help)
            usage
            ;;
        *)
            echo "未知参数: $1"
            usage
            ;;
    esac
done

# 检查必需参数
if [ -z "$PT_PATH" ]; then
    echo "错误: 必须指定 --pt_path 参数"
    usage
fi

# 构建Python命令
PYTHON_CMD="python ${PROJECT_DIR}/export_model.py"
PYTHON_CMD="${PYTHON_CMD} --pt_path ${PT_PATH}"
PYTHON_CMD="${PYTHON_CMD} --imgsz ${IMGSZ}"
PYTHON_CMD="${PYTHON_CMD} --device ${DEVICE}"
PYTHON_CMD="${PYTHON_CMD} --batch ${BATCH}"
PYTHON_CMD="${PYTHON_CMD} --workspace ${WORKSPACE}"
PYTHON_CMD="${PYTHON_CMD} --conf ${CONF}"
PYTHON_CMD="${PYTHON_CMD} --iou ${IOU}"
PYTHON_CMD="${PYTHON_CMD} --max-det ${MAX_DET}"

if [ "$HALF" = false ]; then
    PYTHON_CMD="${PYTHON_CMD} --no-half"
fi

if [ "$NMS" = true ]; then
    PYTHON_CMD="${PYTHON_CMD} --nms"
fi

if [ -n "$DATA" ]; then
    PYTHON_CMD="${PYTHON_CMD} --data ${DATA}"
fi

if [ "$DYNAMIC" = true ]; then
    PYTHON_CMD="${PYTHON_CMD} --dynamic"
fi

if [ "$INT8" = true ]; then
    PYTHON_CMD="${PYTHON_CMD} --int8"
    PYTHON_CMD="${PYTHON_CMD} --fraction ${FRACTION}"
fi

if [ "$QUIET" = true ]; then
    PYTHON_CMD="${PYTHON_CMD} --quiet"
fi

# 打印配置信息
if [ "$QUIET" = false ]; then
    echo "================================"
    echo "模型导出配置"
    echo "================================"
    echo "Conda环境: dl_env"
    echo "项目目录: ${PROJECT_DIR}"
    echo "模型文件: ${PT_PATH}"
    echo "图像尺寸: ${IMGSZ}"
    echo "批处理: ${BATCH}"
    echo "设备: GPU ${DEVICE}"
    echo "半精度: ${HALF}"
    echo "动态形状: ${DYNAMIC}"
    echo "INT8量化: ${INT8}"
    echo "融合NMS: ${NMS}"
    echo "工作空间: ${WORKSPACE}GiB"
    if [ "$NMS" = true ]; then
        echo "NMS置信度: ${CONF}"
        echo "NMS IOU: ${IOU}"
        echo "最大检测: ${MAX_DET}"
    fi
    echo "================================"
    echo ""
fi

# 执行导出命令
echo "开始导出模型..."
eval ${PYTHON_CMD}

# 检查执行结果
if [ $? -eq 0 ]; then
    echo ""
    echo "✅ 导出完成！"
else
    echo ""
    echo "❌ 导出失败！"
    exit 1
fi
